#include "crypto_helper.hpp"
#include "intel_itt_wrapper.hpp"

static inline bool is_option_invalid(const std::string &opt) {
  return opt.empty() || opt[0] != '-';
}

static inline bool stobool(const std::string &str) {
  std::string lower_str;
  for (auto ch : str) {
    lower_str += static_cast<char>(std::tolower(ch));
  }
  //! No need for error detection
  return lower_str == "true";
}

static auto defualt_print_usage = []() {
  // clang-format off
  std::cout << "  -custom1                Enable custom1 option" << std::endl;
  std::cout << "  -custom2 <value>        Set custom2 to the specified value (default: default_custom2_value)" << std::endl;
  // clang-format on
  std::flush(std::cout);
};

static auto bgv_params_print_usage = []() {
  // clang-format off
  std::cout << "Usage: program_name [options]" << std::endl;
  std::cout << "Options:" << std::endl;
  std::cout << "  -hwt <value>            Set the hamming weight of the secret key (default: 120)" << std::endl;
  std::cout << "  -mlb <value>            Set the lower bound of m (default: 20000)" << std::endl;
  std::cout << "  -mub <value>            Set the upper bound of m (default: 30000)" << std::endl;
  std::cout << "  -qblb <value>           Set the minimum bit length of ciphertest modulus q (default: 400)" << std::endl;
  std::cout << "  -c <value>              Set # of columns of Key-Switching matrix (default: 2)" << std::endl;
  std::cout << "  -tb <value>             Set the bit length of plaintext modulus t (default: 30, i.e., 2^30<= t < 2^31)" << std::endl;
  std::cout << "  -custom1                Enable custom1 option" << std::endl;
  std::cout << "  -custom2 <value>        Set custom2 to the specified value (default: default_custom2_value)" << std::endl;
  // clang-format on
  defualt_print_usage();
};

static auto ckks_params_print_usage = []() {
  // clang-format off
  std::cout << "Usage: program_name [options]" << std::endl;
  std::cout << "Options:" << std::endl;
  std::cout << "  -bits <value>           Set the minimum bit length of ciphertest modulus q (default: 119)" << std::endl;
  std::cout << "  -pc <value>             Set the number of bits of precision when data is encoded, encrypted, or decrypted (default: 20)" << std::endl;
  std::cout << "  -custom1                Enable custom1 option" << std::endl;
  std::cout << "  -custom2 <value>        Set custom2 to the specified value (default: default_custom2_value)" << std::endl;
  // clang-format on
  defualt_print_usage();
};

command_line_parser::command_line_parser() {
  print_usage = defualt_print_usage;
}

command_line_parser::command_line_parser(params<helib::BGV> &param) {
  // Preset parameters
  // clang-format off
  add_option("-hwt", std::to_string(param.hwt), "hwt", &param.hwt);
  add_option("-mlb", std::to_string(param.m_lower_bound), "m_lower_bound", &param.m_lower_bound);
  add_option("-mub", std::to_string(param.m_upper_bound), "m_upper_bound", &param.m_upper_bound);
  add_option("-qblb", std::to_string(param.modulus_chain_bits), "modulus_chain_bits", &param.modulus_chain_bits);
  add_option("-c", std::to_string(param.ksm_ncol), "ksm_ncol", &param.ksm_ncol);
  add_option("-tb", std::to_string(param.t_bits), "t_bits", &param.t_bits);
  // clang-format on
  print_usage = bgv_params_print_usage;
}

command_line_parser::command_line_parser(params<helib::CKKS> &param) {
  // Preset parameters
  // clang-format off
  add_option("-mlb", std::to_string(param.m_lower_bound), "m_lower_bound", &param.m_lower_bound);
  add_option("-bits", std::to_string(param.bits), "bits(slight greater than capacity)", &param.bits);
  add_option("-pc", std::to_string(param.precision), "precision", &param.precision);
  // clang-format on
  print_usage = ckks_params_print_usage;
}

int command_line_parser::parse(int argc, char *argv[], bool report) {
  // Parse command line
  for (int i = 1; i < argc; ++i) {
    std::string arg = argv[i];
    if (options_.find(arg) != options_.end()) {
      if (i + 1 < argc && argv[i + 1][0] != '-') {
        // No leading `-` means an option with parameters
        options_[arg] = argv[++i];
      } else {
        // Enalbe option `arg` and set it to "true"
        options_[arg] = "true";
      }
    } else {
      // UnKnown options
      print_usage();
      return 1;
    }
  }
  update(report);
  return 0;
}

void command_line_parser::add_option(const std::string &opt,
                                     const std::string &default_str,
                                     const std::string &name, void *var) {
  ASSERT_PRINTF(!is_option_invalid(opt), "Invalid option: %s\n", opt.c_str());
  options_[opt] = default_str;
  option_info_[opt] = {name, var, false};
}

void command_line_parser::add_option(const std::string &opt,
                                     const std::string &name, void *var) {
  ASSERT_PRINTF(!is_option_invalid(opt), "Invalid option: %s\n", opt.c_str());
  options_[opt] = "false";
  option_info_[opt] = {name, var, true};
}

void command_line_parser::update(bool verbose) {
  if (verbose) {
    std::cout << "============ Command Line Parameters Report ===========\n";
  }
  for (auto it = option_info_.begin(); it != option_info_.end(); it++) {
    const std::string arg = it->first;
    const auto &info = it->second;
    const auto val = options_[arg];
    if (info.is_bool) {
      *static_cast<bool *>(info.var) = stobool(val);
      if (verbose) {
        std::cout << "* " << info.name << " = "
                  << (*static_cast<bool *>(info.var) ? "true" : "false")
                  << std::endl;
      }
    } else {
      *static_cast<unsigned long *>(info.var) = std::stoul(val);
      if (verbose) {
        std::cout << "* " << info.name << " = "
                  << *static_cast<unsigned long *>(info.var) << std::endl;
      }
    }
  }

  if (verbose) {
    std::cout << "=======================================================\n";
  }
}

hmm_crypto_agent::hmm_crypto_agent(std::istream &str) {
  build_context_time_ -= omp_get_wtime();
  context_ = std::shared_ptr<helib::Context>(helib::Context::readPtrFrom(str));
  build_context_time_ += omp_get_wtime();
}

void hmm_crypto_agent::report_context(std::ostream &os) const {
  // Report current context
  os << "=================== Context Report ====================\n";
  context_->printout(os);
  os << "=======================================================\n";
}

// Homomorphic Cryptograph Client
template <typename Scheme>
hmm_crypto_client<Scheme>::hmm_crypto_client(const helib::Context &context) {
  INTEL_ITT_DOMAIN_CREATE(hyper_hmme_domain, "src.hypercube.hmm_crypto_client");
  INTEL_ITT_STRING_HANDLE_CREATE(keygen_handle, "key generation");

  INTEL_ITT_TASK_BEGIN(hyper_hmme_domain, keygen_handle); //! ITT - Begin

  key_gen_time_ -= omp_get_wtime(); //! KeyGen - Start
  // Create a secret key associated with the context
  sk_ = new helib::SecKey(context);
  // Generate the secret key
  sk_->GenSecKey();
  // Compute key-switching matrices that we need
  helib::addSome1DMatrices(*sk_);
  key_gen_time_ += omp_get_wtime(); //! KeyGen - Stop

  INTEL_ITT_TASK_END(hyper_hmme_domain); //! ITT - End
}

template <typename Scheme>
helib::Ctxt hmm_crypto_client<Scheme>::get_empty_ctxt() const {
  return std::move(helib::Ctxt(*sk_));
}

template <typename Scheme>
void hmm_crypto_client<Scheme>::encrypt(helib::Ctxt &ctxt,
                                        const helib::Ptxt<Scheme> &ptxt,
                                        ENC type) {
  // sanity check
  ASSERT_PRINTF(ctxt.getPubKey() == public_key(),
                "Encrypt: public key mismatch\n");
  switch (type) {
  case ENC::PUBKEY: {
    // Get public key part from the secret key
    const helib::PubKey &pk = *sk_;
    // Encrypt using public key
    pk.Encrypt(ctxt, ptxt);
  } break;
  case ENC::SECKEY: {
    // FIXME: No difference?
    helib::EncodedPtxt eptxt;
    ptxt.encode(eptxt);
    sk_->Encrypt(ctxt, eptxt);
  } break;
  default:
    ERROR_PRINTF("Invalid encryption type\n");
    break;
  }
}

template <typename Scheme>
void hmm_crypto_client<Scheme>::decrypt(helib::Ptxt<Scheme> &ptxt,
                                        const helib::Ctxt &ctxt) {
  // Decrypt using secret key
  sk_->Decrypt(ptxt, ctxt);
}

// //* template specialization
// for BGV
template class hmm_crypto_client<helib::BGV>;
// for CKKS
template class hmm_crypto_client<helib::CKKS>;